# -*- coding: utf-8 -*-

import os
import re
from tkinter import *
from tkinter import scrolledtext
from tkinter.filedialog import askopenfilename


class App(Frame):
    #构造函数
    def __init__(self, master=None):
        Frame.__init__(self, master)
        self.MODE = IntVar()
        self.MODE.set(1)
        self.grid()
        self.createWidgets()
        self.parserEntry.insert(0, "。")
        self.rangeSelector.insert(0, "0:0")

    #浏览按钮动作定义
    def browse(self):
        Filelist = askopenfilename()
        self.filterSelector.delete(0, 'end')
        self.filterSelector.insert(0, Filelist)

    #初始化控件
    def createWidgets(self):
        self.modeLabel = Label(self, text="请选择待检索文件类型：")
        self.modeLabel.grid(row=0, sticky=E)
        self.modeSelectorIPO = Radiobutton(
            self, text="招股说明书", value=1, variable=self.MODE)
        self.modeSelectorIPO.grid(row=0, column=1)
        self.modeSelectorAR = Radiobutton(
            self, text="年报", value=2, variable=self.MODE)
        self.modeSelectorAR.grid(row=0, column=2)
        self.nameLabel = Label(self, text="请输入检索方案名称：")
        self.nameLabel.grid(row=1, sticky=E)
        self.parserLabel = Label(self, text="请输入断句标点，默认为句号：")
        self.parserLabel.grid(row=2, sticky=E)
        self.ruleLabel = Label(self, text="请输入检索规则：")
        self.ruleLabel.grid(row=3, sticky=E)
        self.nameEntry = Entry(self, width=50)
        self.nameEntry.grid(row=1, column=1, columnspan=3)
        self.parserEntry = Entry(self, width=50)
        self.parserEntry.grid(row=2, column=1, columnspan=3)
        self.ruleEntry = Text(self, height=10, width=50)
        self.ruleEntry.grid(row=3, column=1, columnspan=3)
        self.rangeLabel = Label(self, text="请设置前后步长，以:分隔：")
        self.rangeLabel.grid(row=4, sticky=E)
        self.rangeSelector = Entry(self, width=50)
        self.rangeSelector.grid(row=4, column=1, columnspan=3)
        self.fileterLabel = Label(self, text="请选择要过滤的文件名列表：")
        self.fileterLabel.grid(row=5, sticky=E)
        self.filterSelector = Entry(self, width=44)
        self.filterSelector.grid(row=5, column=1, columnspan=2)
        self.browseButton = Button(self, text="浏览", command=self.browse)
        self.browseButton.grid(row=5, column=3)
        self.button = Button(self, text="开始检索", command=self.main)
        self.button.grid(row=6, columnspan=4, sticky=E)
        self.consoleText = scrolledtext.ScrolledText(self)
        self.consoleText.grid(row=7, columnspan=4)
        self.consoleText.config(state=DISABLED)

    def printInText(self, text):
        self.consoleText.config(state=NORMAL)
        self.consoleText.insert(END, text)
        self.consoleText.config(state=DISABLED)
        self.consoleText.see(END)
        self.update()

    #定义检索函数1&2&3
    def extract(self, text, regex):  #text为分割后的列表格式文本，rule为|连接的单条检索式
        results = []
        n_list = []
        for n in range(len(text)):
            match = regex.search(text[n])
            if match:
                n_list.append(n)
        for i in n_list:
            start = max(0, i - self.backward)
            end = min(i + 1 + self.forward, len(text) - 1)
            result = "//".join(text[start:end])
            result = result.replace(",", "")
            result = result.replace("\n", "/")
            results.append(result)
        return results

    def extract2(self, text, regex):
        results = []
        n_list = []
        for n in range(len(text)):
            match = regex.search(text[n])
            if match:
                n_list.append(n)
        for i in n_list:
            result = text[i]
            results.append(result)
        return results

    def not_extract(self, text, regex):
        results = []
        n_list = []
        for n in range(len(text)):
            match = regex.search(text[n])
            if not match:
                n_list.append(n)
        for i in n_list:
            result = text[i]
            results.append(result)
        return results

    #主函数
    def main(self):

        #获得环境变量
        workpath = os.path.dirname(__file__)
        path_input = os.path.join(workpath, "input")
        output_name = self.nameEntry.get()
        output_file = output_name + ".csv"
        path_output = os.path.join(workpath, "output")

        #初始化输出文件
        with open(
                os.path.join(path_output, output_file), "w",
                encoding="utf-8") as f:
            if self.MODE.get() == 1:
                f.write('code,id,content\n')
            else:
                f.write('code,year,id,content\n')

        #获取检索规则
        punc = self.parserEntry.get()

        rule = self.ruleEntry.get("0.0", "end")
        rule_list = rule.split("\n")[:-1]

        if "NOT" in rule_list[-1]:
            out_rule = rule_list[-1][3:]
            rule_list.pop()
        else:
            out_rule = ""

        ranges = self.rangeSelector.get()
        self.backward = int(ranges.split(":")[0])
        self.forward = int(ranges.split(":")[1])

        #读取待检索文件列表并筛选(如有)
        filterlist = []
        if self.filterSelector.get() != "":
            with open(self.filterSelector.get()) as f:
                filterlist = f.read().split()

        file_list = [
            x for x in os.listdir(path_input) if x.split(".")[-1] == "txt"
        ]
        count = 0

        if len(filterlist) > 0:
            if self.MODE.get() == 1:
                file_list = [x for x in file_list if x[11:17] in filterlist]
            elif self.MODE.get() == 2:
                file_list = [
                    x for x in file_list if x[5:11] + "," + x[:4] in filterlist
                ]
            else:
                raise ValueError

        #主循环
        for each in file_list:

            #读取文件及股票代码, 年份
            if self.MODE.get() == 1:
                code = each[11:17]
            elif self.MODE.get() == 2:
                year = each[:4]
                code = each[5:11]
            else:
                raise ValueError

            try:
                with open(os.path.join(path_input, each), encoding="gbk") as f:
                    source = f.read()
            except:
                with open(
                        os.path.join(path_input, each),
                        encoding="utf-8",
                        errors="ignore") as f:
                    source = f.read()

            #调用检索函数
            text = re.split(r'%s' % punc, source)
            for i in range(len(rule_list)):
                regex = re.compile(r'%s' % rule_list[i])
                if i == 0:
                    text = self.extract(text, regex)
                else:
                    text = self.extract2(text, regex)
            if out_rule != "":
                regex = re.compile(r'%s' % out_rule)
                text = self.not_extract(text, regex)

            #结果输出
            with open(
                    os.path.join(path_output, output_file), "a",
                    encoding="utf-8") as f:
                if len(text) == 0:
                    f.write(code)
                    f.write(',')
                    if self.MODE.get() == 2:
                        f.write(year)
                        f.write(",")
                    f.write(str(0))
                    f.write(',')
                    f.write('无记录')
                    f.write('\n')
                else:
                    for j in range(len(text)):
                        f.write(code)
                        f.write(',')
                        if self.MODE.get() == 2:
                            f.write(year)
                            f.write(",")
                        f.write(str(j + 1))
                        f.write(',')
                        f.write(text[j])
                        f.write('\n')
            count = count + 1
            self.printInText("当前进度: {:.2f}%\n".format(
                count * 100 / len(file_list)))
        self.printInText("方案{}完成！\n已输出至{}\n".format(
            output_name,
            os.path.realpath(os.path.join(path_output, output_file))))


if __name__ == "__main__":
    app = App()
    app.master.title("Extarct 5.0")
    app.mainloop()
